#
# Read annealing_trace and trace.txt from Phyesta.
#
# Plot each BL proportion vs. the summed length of all branches in all BL proportions
#

import os
import re
import sys
import dendropy

from argParser import ArgParser
argp = ArgParser(sys.argv)

atpath = argp.getArg("--annealing_trace")
tpath = argp.getArg("--phyesta_trace")

points_lnl_props = [] # each value is an array [lnl, treelength, [blprops]]

fin = open(atpath, "r")
last_lnl = None
while 1:
    lines = fin.readlines(100000)
    if not lines:
        break
    for l in lines:
        if l.__contains__("proposition"):
            l = l.strip()
            lnl = l.split()[4]
            last_lnl = float(re.sub("\]", "", lnl))
        elif l.__contains__("BL proportions"):
            l = l.strip()
            l = re.sub("\]", "", l)
            l = re.sub("\[", "", l)
            props = l.strip()[3:]
            points_lnl_props.append( [last_lnl, props] )  
            #print [last_lnl, props]
fin.close()

points_lnl_treelen = []
fin = open(tpath, "r")
last_lnl = None
last_treelens = []
while 1:
    lines = fin.readlines(100000)
    if not lines:
        break
    for l in lines:
        l = l.strip()
        if l.startswith("["):
            # we found a new entry.
            if last_lnl != None: # record last data
                points_lnl_treelen.append( [last_lnl, last_treelens] )
                print [last_lnl, last_treelens]
            # reset our lnl and treelens variables:                
            last_lnl = float(re.sub("\[", "", l.split("]")[0]))
            last_treelens = []
        if l.__contains__("("): # this line contains a tree:
            l = re.sub("\[.*\]", "", l)
            tree = dendropy.Tree()
            tree.read_from_string( l, "newick")
            last_treelens.append( tree.length() )
fin.close()


# Adjust for the fact that points_props will (potentially) contain data from the pre-optimization step.
# ...and we want to ignore this pre-opt data.
delta = points_lnl_treelen.__len__() - points_lnl_props.__len__()

# points_lnl_treelen[delta] corresponds to poins_lnl_props[0]

